package com.gitee.melin.bee.core.grpc.error;

import io.grpc.ForwardingServerCallListener;
import io.grpc.ServerCall;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import org.apache.dubbo.rpc.RpcContext;

/** huaixin 2022/1/28 9:53 AM */
public class GrpcExceptionListener<ReqT, RespT>
        extends ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {

    private final GrpcExceptionResponseHandler exceptionHandler;

    private final ServerCall<ReqT, RespT> serverCall;

    private final InetSocketAddress socketAddress;

    /**
     * Creates a new exception handling grpc server call listener.
     *
     * @param delegate The listener to delegate to (Required).
     * @param serverCall The server call to used to send the error responses (Required).
     * @param exceptionHandler The exception handler to use (Required).
     */
    protected GrpcExceptionListener(
            final ServerCall.Listener<ReqT> delegate,
            final ServerCall<ReqT, RespT> serverCall,
            final GrpcExceptionResponseHandler exceptionHandler,
            final SocketAddress socketAddress) {

        super(delegate);
        this.serverCall = serverCall;
        this.exceptionHandler = exceptionHandler;
        this.socketAddress = (InetSocketAddress) socketAddress;
    }

    @Override
    public void onMessage(final ReqT message) {
        try {
            RpcContext.getContext().setRemoteAddress(socketAddress);
            super.onMessage(message);
        } catch (final Throwable error) {
            // For errors thrown in the request's StreamObserver#onNext
            this.exceptionHandler.handleError(this.serverCall, error);
        }
    }

    @Override
    public void onHalfClose() {
        try {
            RpcContext.getContext().setRemoteAddress(socketAddress);
            super.onHalfClose();
        } catch (final Throwable error) {
            // For errors from unary grpc method implementation methods directly (Not via
            // StreamObserver)
            // For errors thrown in the request's StreamObserver#onCompleted
            this.exceptionHandler.handleError(this.serverCall, error);
        }
    }
}
